import datetime

from langchain.tools import tool, StructuredTool
from langchain_openai import ChatOpenAI

from config.load_key import load_key

llm = ChatOpenAI(
    model="Qwen/Qwen2.5-7B-Instruct",
    api_key=load_key("siliconflow_api_key"),
    base_url=load_key("siliconflow_base_url"),
)


# 定义工具 注意要添加注解 方法中通过""""""注释告诉大模型此方法的作用
# 自定义工具名称：@tool注解是LangChain官方提供的一种装饰器，用于定义工具。他接受一个参数，这个参数就是工具名称。所以我们可以在声明工具时，自己定义工具的名字。如果不指定，默认就是方法名。
# 自定义工具描述：在定义工具时，需要自定义工具描述。这个描述是给大模型用的，大模型会根据这个描述来判断是否需要调用这个工具。描述信息可以在方法中直接添加注释，也可以在@tool注解的desciption属性中定制。所以定义工具方法时，最好把注释写清楚。另外，在定义工具时，除了需要定义工具的描述，还可以定义参数的描述，这样大模型也能根据参数的描述来判断如何调用这个工具。
@tool("get_current_date")
def get_current_date():
    """获取今天日期"""
    return datetime.datetime.today().strftime("%Y-%m-%d")


@tool(description="获取某个城市的天气")
def get_city_weather(city: str):
    """获取某个城市的天气
    Args:
        city: 具体城市
    """
    return "城市" + city + "，今天天气不错"


# 除了使用@tool注解外，langchain中还提供了结构化工具StructuredTool.from_function同样可以用来定制工具。这种方式比@tool具有更多的可配置性，且不需要太多的代码
def bad_weather_tool(city: str):
    """获取某个城市的天气
    Args:
        city: 具体城市
    """
    return "城市" + city + "，今天天气不太好"


# 定义工具。这个方法中有更多参数可以定制
weatherTool = StructuredTool.from_function(func=bad_weather_tool, description="获取某个城市的天气", name="bad_weather_tool")

# 大模型绑定工具
llm_with_tools = llm.bind_tools([get_current_date, weatherTool])
# 工具容器
all_tools = {
    "get_current_date": get_current_date,
    "bad_weather_tool": weatherTool,
    "get_city_weather": get_city_weather
}
# 把所有消息存到一起
query = "今天是几月几号，北京天气怎么样"
messages = [query]
# 询问大模型，大模型会判断需要调用工具，并返回一个工具调用请求
ai_msg = llm_with_tools.invoke(messages)
print(ai_msg)
messages.append(ai_msg)
# 打印需要调用的工具
print(ai_msg.tool_calls)
if ai_msg.tool_calls:
    for tool_call in ai_msg.tool_calls:
        selected_tool = all_tools[tool_call["name"].lower()]
        tool_msg = selected_tool.invoke(tool_call)
        messages.append(tool_msg)
print(llm_with_tools.invoke(messages).content)
